#这个代码会对图像进行一个高频增强的操作，帮助我看懂高频是什么东西
import numpy as np
import cv2
import matplotlib.pyplot as plt

def frequency_masks(shape, low_ratio=0.1, high_ratio=0.3):
    H, W = shape
    center_u, center_v = H // 2, W // 2
    U, V = np.ogrid[:H, :W]
    dist = np.sqrt((U - center_u)**2 + (V - center_v)**2)
    max_dist = dist.max()
    low_mask = dist <= low_ratio * max_dist
    mid_mask = (dist > low_ratio * max_dist) & (dist <= high_ratio * max_dist)
    high_mask = dist > high_ratio * max_dist
    return low_mask, mid_mask, high_mask

def normalize_to_uint8(img):
    img = img - img.min()
    img = img / (img.max() + 1e-8)
    return (img * 255).astype(np.uint8)

def process_high_freq_enhancement(img_path, alpha=2.5):
    img_bgr = cv2.imread(img_path)
    img_ycc = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCrCb)
    y, cr, cb = cv2.split(img_ycc)
    y = y.astype(np.float32)

    H, W = y.shape
    low_mask, mid_mask, high_mask = frequency_masks((H, W))

    # 频谱处理
    F = np.fft.fft2(y)
    F_shifted = np.fft.fftshift(F)
    F_high = F_shifted * high_mask
    y_high = np.fft.ifft2(np.fft.ifftshift(F_high)).real

    # 高频增强：原图 + α * 高频
    y_enhanced = y + alpha * y_high
    y_enhanced = np.clip(y_enhanced, 0, 255).astype(np.uint8)

    # 合成增强图像
    img_ycc_enhanced = cv2.merge([y_enhanced, cr, cb])
    img_rgb_enhanced = cv2.cvtColor(img_ycc_enhanced, cv2.COLOR_YCrCb2RGB)

    # 显示结果
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.figure(figsize=(12, 5))
    plt.subplot(1, 3, 1)
    plt.imshow(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
    plt.title("原图"); plt.axis('off')
    plt.subplot(1, 3, 2)
    plt.imshow(normalize_to_uint8(y_high), cmap='gray')
    plt.title("高频部分 (增强前)"); plt.axis('off')
    plt.subplot(1, 3, 3)
    plt.imshow(img_rgb_enhanced)
    plt.title(f"高频增强图 (α={alpha})"); plt.axis('off')
    plt.tight_layout()
    plt.show()

# 执行路径替换为你自己的图像路径
process_high_freq_enhancement('D:\\2D-DFT\\image\\testhaze3.png', alpha=3)
